#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time     : 2024/6/25 9:41
# @Author   : 22759
# @Email    : zjh13092921608@163.com
# @Project  : S2
# @File     : batchViews.py
# @Software : PyCharm
import os
import re
import traceback
from typing import Optional, Dict, Tuple, List
from concurrent.futures import ThreadPoolExecutor

from django.db.models import Q
from django.conf import settings
from django.shortcuts import render
from django.http import HttpRequest
from django.views.decorators.csrf import csrf_exempt

from celery_tasks.tasks import trigger_batch
from utils.decorators import service_required
from utils.viewutil import savelog, ajaxresponse, pages, get_user
from apps.assetmgnt.models import DevType, Device, BatchOperation


class BatchOptions:

    def __init__(self, request, behavior):
        self.request:   HttpRequest = request
        self.behavior:  Optional[int] = behavior
        self.pattern:   Optional[str]
        self.statement: Optional[str]
        self.status_choices:   Optional[Dict] = dict(BatchOperation.STATUS_CHOICES)
        self.behavior_choices: Optional[Dict] = dict(BatchOperation.BEHAVIOR_CHOICES)
        self.upgrade_devices:  Optional[Tuple] = ('T3-3.1', 'T3-5.0', 'A4-3.0', 'LAS')
        self.upgrade_library:  Optional[Tuple] = ('T3-3.1', 'T3-5.0', 'A4-3.0')
        self.exclude_devices:  Optional[Tuple] = (
            'fortress_machine', 'fortress_machine_V2',  'fortress_machine_base', 'next_filrewall')
        self.device_types:     Optional[List] = self._device_types()
        """存放文件位置默认设置"""
        if not os.path.exists(settings.BATCH_PATH):
            os.makedirs(settings.BATCH_PATH, exist_ok=True)

    def batch_infomation(self):
        """批量操作基本信息"""
        devices_type = self.device_types
        devices_dict = self._devices()
        running = BatchOperation.objects.filter(Q(status='0') | Q(status__isnull=True)).exists()
        batch_objects = BatchOperation.objects.filter(behavior=self.behavior).all()
        for batch_object in batch_objects:
            batch_object.cause = batch_object.cause if batch_object.cause else ''
            batch_object.status_text = self.status_choices.get(batch_object.status, u'未启动')
        s, e, pagination = pages(self.request, batch_objects)
        return {
            'running': running,
            'devicesType': devices_type,
            'devicesDict': devices_dict,
            'datas': batch_objects[s:e],
            'pagination': pagination,
        }

    def batch_operation(self):
        """下发批量操作任务"""
        """定义标识符"""
        head, message = False, '执行失败'
        """请求参数"""
        files = self.request.FILES.get('files')
        filename = self.request.POST.get('filename')
        start = int(self.request.POST.get('start', 0))
        device_type = self.request.POST.get('deviceType')
        device_ides = self.request.POST.getlist('deviceIdes')
        end = self.request.POST.get('end', 'false') == 'true'
        """校验内容"""
        if not device_ides or not files:
            head, message = False, '缺少必要参数'
        else:
            """获取请求类型"""
            dev_type = DevType.objects.filter(id=device_type).first()
            """文件校验"""
            if not self._check_file(filename, dev_type.ename if self.behavior == 1 else 'library'):
                head, message = False, f'{filename}无法进行升级'
            else:
                self._save_file(dev_type.ename, start)
                if end:
                    """保存配置"""
                    batch_ides = self._save_batch(device_ides)
                    """调用触发任务"""
                    trigger_batch.delay(dev_type.ename, device_ides, self.behavior, batch_ides)
                    head, message = True, '批量任务已下发'
        return head, message

    def _check_file(self, filename, device_type):
        """校验文件函数"""
        if 'T3' in device_type:
            """校验防火墙升级文件"""
            # self.pattern = re.compile(r'-ext4.img.gz$')
            self.pattern = re.compile(r'.bin$')  # ARM 64使用bin升级
        elif any(mark in device_type for mark in {'A4', 'LAS'}):
            """校验流量审计升级文件"""
            self.pattern = re.compile(r'-update.tar.gz$')
        elif 'library' == device_type:
            """校验特征库升级文件"""
            self.pattern = re.compile(r'.gpg$')
        if self.pattern.search(filename):
            return True
        return False

    def _save_file(self, device_type, start):
        """保存升级文件"""
        filename = ''
        if self.behavior == 1:
            if any(mark in device_type for mark in ('A4', 'LAS')):
                filename = 'upgrade_system-update.tar.gz'
            elif 'T3' in device_type:
                filename = 'upgrade_system.bin'
        elif self.behavior == 2:
            filename = 'upgrade_library.gpg'
        if filename:
            with ThreadPoolExecutor(max_workers=4) as executor:
                future = executor.submit(self._write_file, f'{settings.BATCH_PATH}/{filename}', self.request.FILES['files'], start)
                future.result()
        return True

    def _save_batch(self, devices, content=None):
        """保存批量操作记录"""
        result, user = list(), get_user(self.request)
        try:
            for device in devices:
                self.statement = f"""
                SELECT
                    t_device.id,
                    t_device.name,
                    t_device.ip,
                    t_devtype.name as type
                FROM
                    t_device
                    LEFT JOIN t_devtype ON t_device.dtype_id = t_devtype.id
                WHERE
                    t_device.id = {device};
                """
                device_objs = Device.objects.raw(self.statement)
                if device_objs:
                    device_obj = device_objs[0]
                    row = {
                        'name': device_obj.name,
                        'ip': device_obj.ip,
                        'type': device_obj.type,
                        'behavior': self.behavior,
                    }
                    batch_obj = BatchOperation.objects.create(**row)
                    result.append(batch_obj.id)
            savelog(
                request=self.request,
                type='sys_config',
                subject=u'系统',
                object=self.behavior_choices[self.behavior],
                priority=6,
                module='batch_operation',
                message=u'{user}设置批量{behavior}任务成功'.format(**{
                    'user': user.username,
                    'behavior': self.behavior_choices[self.behavior]
                })
            )
            return result
        except Exception as error:
            savelog(
                request=self.request,
                type='sys_config',
                subject=u'系统',
                object=self.behavior_choices[self.behavior],
                priority=6,
                module='batch_operation',
                message=u'{user}下发批量{behavior}任务失败，失败原因：{error}'.format(**{
                    'user': user.username,
                    'behavior': self.behavior_choices[self.behavior],
                    'error': str(error)
                })
            )
            return []

    def _device_types(self):
        """可升级设备类型"""
        ename = self.upgrade_library if self.behavior == 2 else self.upgrade_devices
        upgrade_types = DevType.objects.filter(ename__in=ename).all()
        return [{'id': types.id, 'name': types.name} for types in upgrade_types]

    def _devices(self):
        """可升级设备"""
        devices_dict, devices_mark = dict(), set()
        devices_lists = Device.objects.filter(dtype_id__in=tuple(types['id'] for types in self.device_types)).all()
        for device in devices_lists:
            if device.dtype_id not in devices_mark:
                devices_mark.add(device.dtype_id)
                devices_dict[device.dtype_id] = [{'id': device.id, 'name': device.name}]
            else:
                devices_dict[device.dtype_id].append({'id': device.id, 'name': device.name})
        return devices_dict

    def _write_file(self, file_path, chunks, start):
        with open(file_path, 'ab') as file:
            file.seek(start)
            for chunk in chunks:
                file.write(chunk)


@csrf_exempt
@service_required('sys_config')
def batch(request):
    """批量操作设备中枢"""
    return render(request, 'assettpts/batchControls.html', {})


@csrf_exempt
@service_required('sys_config')
def batch_upgrade(request):
    """批量升级管理设备"""
    """调用批量升级设置类"""
    batch_options = BatchOptions(request=request, behavior=1)
    if request.method == 'POST':
        """批量升级系统任务触发"""
        head, message = batch_options.batch_operation()
        return ajaxresponse({'head': 'OK' if head else '', 'body': message, 'redirect': '/assetmgnt/batch/'})
    """批量升级系统信息"""
    context = batch_options.batch_infomation()
    return render(request, template_name='batchUpdateSystem.html', context=context)


@csrf_exempt
@service_required('sys_config')
def batch_libaray(request):
    """批量升级特征库"""
    """定义标识符"""
    head, message = False, str('执行失败')
    devices_dict, devices_mark = dict(), set()
    """调用批量升级设置类"""
    batch_options = BatchOptions(request=request, behavior=2)
    if request.method == 'POST':
        """批量升级特征库任务触发"""
        head, message = batch_options.batch_operation()
        return ajaxresponse({'head': 'OK' if head else '', 'body': message, 'redirect': '/assetmgnt/batch/'})
    """批量升级特征库信息"""
    context = batch_options.batch_infomation()
    return render(request, template_name='batchUpdateLibrary.html', context=context)


@csrf_exempt
@service_required('sys_config')
def batch_reboot(request):
    """批量重启"""
    """调用批量升级设置类"""
    batch_options = BatchOptions(request=request, behavior=3)
    '''获取可控重启设备'''
    exclude_types = DevType.objects.filter(ename__in=batch_options.exclude_devices).values_list('id', flat=True)
    exclude_types = tuple(type_id for type_id in exclude_types)
    devices = Device.objects.exclude(dtype_id__in=exclude_types).all()
    return render(request, template_name='batchReboot.html', context={'devices': devices})
